# -*- coding: utf-8 -*-
# @Time    : 2022/12/23 15:22
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : gen_particles.py
# @Software: PyCharm

import numpy
import pandas
import scipy.constants as C

# LGB1
# inj_angle = 0.023209725  # rad
# LGB2
inj_angle = 0.038817 # rad
# LGB3
# inj_angle = 0.025083#0.012819 # rad

E0 = 2.2e9  # 5e6  # 2.2e9
G = C.e * E0 / (C.m_e * C.c ** 2)
beta = (1 - 1 / G ** 2) ** .5
GB = G * beta
x_high_mm = 13  # 20
z_magnet_start = 0  # -279.964879e-3 # -52.53e-3# 磁铁端部位置
xs_ref = numpy.arange(-x_high_mm, x_high_mm + 1, 1) * 1e-3  # @ z = z_magnet_start
df = pandas.DataFrame(columns='x y z GBx GBy GBz nmacro rmacro m q'.split())
z0 = -400e-3  # -400e-3  # -150e-3# z_magnet_start  #-45e-3  #-150e-3  # 粒子实际入射位置
# setstartpar("beam",z0*tan(injection_angle),0,z0, GB*sin(injection_angle),0,GB*cos(injection_angle),me, -|qe|, 1) ;
x_center_par = (z0 - z_magnet_start) * numpy.tan(inj_angle)
df['x'] = x_center_par + xs_ref
df['y'] = 0
df['z'] = z0
df['GBx'] = GB * numpy.sin(inj_angle)
df['GBy'] = 0
df['GBz'] = GB * numpy.cos(inj_angle)
df['q'] = -C.e
df['m'] = C.m_e
df['nmacro'] = 1
df['rmacro'] = 1e-9
df.to_csv('e.txt', index=False, sep='\t')
index = df['x'] == x_center_par
df.loc[index, :].to_csv('e_ideal.txt', index=False, sep='\t')
print("Index of center point: %d" % numpy.arange(len(df))[index])
